from function import *
import math
from Loss import Loss
from config import *

class Adam(Loss):
    def __init__(self,fName,config):
        self.fName = fName
        self.dfName = 'df_' + fName
        self.config = config
    def minimize(self):
        z_history = []
        x1_history = []
        x2_history = []
        x1,x2,lr = self.config.x0[0],self.config.x0[1],self.config.lr
        beta = 0.5
        epochs = self.config.epochs
        beta2 = 0.5
        z = eval(self.fName)(x1, x2)
        z_history.append(z)
        s1 ,s2 = 0,0
        v1, v2 = 0, 0
        for epoch in range(epochs):
            dx1, dx2 = eval(self.dfName)(x1, x2)
            v1 = beta * v1 + (1 - beta) * dx1
            v2 = beta * v2 + (1 - beta) * dx2

            s1 = beta2*s1+(1-beta2)* dx1 ** 2
            s2 = beta2*s2+(1-beta2)* dx2 ** 2

            lr1,lr2 = lr/ math.sqrt(s1),lr/ math.sqrt(s2)
            x1 -= lr1  * v1
            x2 -= lr2  * v2

            x1_history.append(x1)
            x2_history.append(x2)
            z = eval(self.fName)(x1, x2)
            z_history.append(z)

            wandb.log({
                'loss': z,
            })
        return  z_history,x1_history,x2_history

def main():
    wandb.init()
    optimizer = Adam(fName,wandb.config)
    # optimizer = getattr(optim, wandb.config.optimizer)(fName, wandb.config)
    z_history,x1_history,x2_history = optimizer.minimize()
    plotAll(fName,np.array(x1_history), np.array(x2_history))
if __name__ == '__main__':
    sweep_id = wandb.sweep(sweep=sweep_configuration, project="优化算法")
    wandb.agent(sweep_id, function=main, count=1)